#pragma once

#include <string>
#include <functional>
#include <memory>
#include <unordered_map>
#include <mutex>

#include "socket.hpp"
#include "thread_pool.hpp"

#include "log.hpp"
extern Log log;

class TCP_server
{
public:
    typedef std::function<void(std::string&, std::string&)> func_t;
    typedef TCP_server self;
public:
    TCP_server()
        : _is_running(false)
    {}
private:
    TCP_server(const self&) = delete;
    self& operator=(const self&) = delete;
private:
    static void tcp_task(int sockfd, bool non_blocking, const std::string client_ip, uint16_t client_port, const func_t& task_func)
    {
        auto context = std::make_shared<ConnectionContext>(sockfd, client_ip, client_port);
        {
            std::lock_guard<std::mutex> lock(_mtx);
            _connections[sockfd] = context;
        }
        
        try
        {
            while (true)
            {
                constexpr size_t recv_size = 1024;
                char recv_buffer[recv_size] = { 0 };
                ssize_t n = recv(sockfd, recv_buffer, sizeof(recv_buffer) - 1, 0);
                if (n > 0)
                {
                    context->buffer.append(recv_buffer, n);
                    
                    std::string outbuffer;
                    task_func(context->buffer, outbuffer);

                    if (!outbuffer.empty())
                    {
                        n = send(sockfd, outbuffer.c_str(), outbuffer.size(), 0);
                        if (n < 0)
                        {
                            log(Log::Warning, "write error, sockfd: %d, clientip: %s, clientport: %d, errstring: %s", sockfd, client_ip.c_str(), client_port, strerror(errno));
                            break;
                        }
                    }
                }
                else if (n == 0)
                {
                    log(Log::Info, "%s:%d quit, server close sockfd: %d", client_ip.c_str(), client_port, sockfd);
                    break;
                }
                else
                {
                    // 非阻塞模式下处理EAGAIN或EWOULDBLOCK
                    if (non_blocking && (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR))
                        continue;
                    
                    log(Log::Warning, "read error, sockfd: %d, clientip: %s, clientport: %d, errstring: %s", sockfd, client_ip.c_str(), client_port, strerror(errno));
                    break;
                }
            }
        }
        catch (const std::exception& e)
        {
            log(Log::Warning, "exception in connection %d: %s", sockfd, e.what());
        }
        
        // 连接关闭后清理资源
        close(sockfd);
        {
            std::lock_guard<std::mutex> lock(_mtx);
            _connections.erase(sockfd);
        }
    }
public:
    void init(uint16_t port)
    {
        _listen_socket.Create(SOCK_STREAM);
        _listen_socket.Bind(port);
        _listen_socket.Listen();
        _thread_pool.start();
    }

    void run(const func_t& func, int non_blocking_flag = 0)
    {
        bool non_blocking = (non_blocking_flag == 1);
        
        if (_is_running == false)
        {
            _is_running = true;
            while (_is_running)
            {
                std::string client_ip;
                uint16_t client_port;
                int sockfd = _listen_socket.Accept(client_ip, client_port);
                if (sockfd < 0)
                    continue;
                
                // 设置非阻塞模式
                if (non_blocking)
                {
                    int flags = fcntl(sockfd, F_GETFL, 0);
                    if (flags != -1)
                        fcntl(sockfd, F_SETFL, flags | O_NONBLOCK);
                    else
                        log(Log::Warning, "non blocking failed sockfd: %d", sockfd);
                }
                
                std::function<void(void)> thread_func = std::bind(tcp_task, sockfd, non_blocking, std::move(client_ip), client_port, func);
                _thread_pool.push(thread_func);
            }
        }
    }

    ~TCP_server()
    {
        _listen_socket.Close();
    }
private:
    Socket _listen_socket;
    bool _is_running;
    
    struct ConnectionContext
    {
        std::string buffer;
        std::string client_ip;
        uint16_t client_port;
        int sockfd;
        
        ConnectionContext(int fd, const std::string& ip, uint16_t port)
            : sockfd(fd), client_ip(ip), client_port(port)
        {}
    };
    
    static std::unordered_map<int, std::shared_ptr<ConnectionContext>> _connections;
    static std::mutex _mtx;

#ifdef __TEST_DEBUG__
    ThreadPool<10> _thread_pool;
#else
    ThreadPool<> _thread_pool;
#endif
};
